from common import pd, np
from scipy.stats import skew, kurtosis
from numpy.typing import NDArray


def generate_data(sample_number: int = 100000, failure_rate: float = 0.06) -> pd.DataFrame:
    """
    生成设备传感器模拟数据和故障日志(a-d四个传感器;0-3四个故障类型;0-正常)
    :param:
        sample_number: 总样本数
        failure_rate: 故障发生比例
    :return:
        DataFrame: 包含时间戳、传感器数据和故障类型
    """
    # 生成时间序列
    timestamps = pd.date_range(start='2025-01-01', periods=sample_number, freq='min')

    # 生成基础传感器数据（正常状态）
    random_a = np.random.normal(50, 5, sample_number)
    random_b = np.random.normal(30, 3, sample_number)
    random_c = np.random.normal(100, 10, sample_number)
    random_d = np.random.normal(80, 8, sample_number)

    # 创建故障 - 0=正常，1=故障-1，2=故障-2，3=故障-3
    failures = np.zeros(sample_number, dtype=int)

    # 随机选择故障点
    failure_indices = np.random.choice(sample_number, int(sample_number * failure_rate), replace=False)

    # 分配故障类型
    failure_types = np.random.choice([1, 2, 3], len(failure_indices), p=[0.5, 0.3, 0.2])

    # 在故障点附近添加异常模式
    for idx, fault_type in zip(failure_indices, failure_types):
        # 标记故障点（故障发生时刻）
        failures[idx] = fault_type

        # 在故障前20分钟开始出现异常（预警期）
        start = max(0, idx - 20)

        # 根据故障类型创建不同的异常模式
        if fault_type == 1:  # 故障-1：传感器a和b异常
            # 传感器a: 均值逐渐升高
            random_a[start:idx] += np.linspace(0, 20, idx - start)
            # 传感器b: 波动增大
            random_b[start:idx] += np.random.normal(0, 7, idx - start)

        elif fault_type == 2:  # 故障-2：传感器c和d异常
            # 传感器c: 出现周期性尖峰
            for i in range(start, idx):
                if i % 6 == 0:  # 每6分钟一个尖峰
                    random_c[i] += np.random.uniform(25, 45)
            # 传感器d: 缓慢下降
            random_d[start:idx] -= np.linspace(0, 15, idx - start)

        elif fault_type == 3:  # 故障-3：所有传感器异常
            # 传感器a: 阶梯式上升
            steps = np.arange(0, idx - start, 6)
            for step in steps:
                random_a[start + int(step):idx] += 5
            # 传感器b: 持续下降
            random_b[start:idx] -= np.linspace(0, 10, idx - start)
            # 传感器c: 随机波动
            random_c[start:idx] += np.random.normal(0, 8, idx - start)
            # 传感器d: 骤降后恢复
            if idx - start > 12:
                random_d[start:start + 6] -= 15
                random_d[start + 6:start + 12] += 10
                random_d[start + 12:idx] -= np.linspace(0, 20, idx - start - 12)

    # 添加随机噪声
    noise_a = np.random.normal(0, 1.5, sample_number)
    noise_b = np.random.normal(0, 1, sample_number)
    noise_c = np.random.normal(0, 2.5, sample_number)
    noise_d = np.random.normal(0, 2, sample_number)

    # 创建DataFrame
    datas = pd.DataFrame({
        'timestamp': timestamps,
        'sensor_a': random_a + noise_a,
        'sensor_b': random_b + noise_b,
        'sensor_c': random_c + noise_c,
        'sensor_d': random_d + noise_d,
        'failure_type': failures
    })
    return datas


def extract_data(datas: pd.DataFrame, label_key: str, feature_keys: tuple[str, ...]) -> tuple[np.ndarray, np.ndarray]:
    """
   将数据提取为特征和标签数据
   :param:
       datas: DataFrame格式的原始数据
   :return:
       tuple[features: 特征数据二维数组, labels: 标签数据一维数组]
   """
    # 多变量时间序列特征
    features = datas[list(feature_keys)].values
    # 分类标签
    labels = datas[label_key].values
    return features, labels


def window_data(datas: np.ndarray, idx=9, window_size=10) -> np.ndarray:
    """
    提取一个窗口的数据,从索引idx位置向前取值, 包含当前idx对应的元素
    :param:
    datas: 要窗口化的数据
    idx: 起始位置索引
    window_size: 窗口大小
    :return: 一个窗口的数据
    """
    index = idx + 1
    if index < window_size:
        raise ValueError("idx is too small, less than one window!")
    rows = len(datas)
    if index > rows:
        raise ValueError("idx out of range!")
    if rows < window_size:
        raise ValueError("Too few data rows, less than one window!")

    return datas[index - window_size: index]


def windowing_data(datas: np.ndarray, window_size=10) -> np.ndarray:
    """
   窗口化数据,将数据按窗口大小滑动拆分
   :param:
       datas: 要窗口化的数据
       window_size: 窗口大小
   :return: 窗口化的数据数组
   """
    windows = []
    rows = len(datas)
    if rows < window_size:
        raise ValueError("Too few data rows, less than one window!")
    for idx in range(window_size - 1, rows):
        windows.append(window_data(datas, idx, window_size))
    return np.array(windows)


def extract_statistical_features(feature_values: np.ndarray) -> pd.DataFrame:
    """
    从原始特征中提取统计特征
    :param
        feature_values: nparray,二维数组 [[1,2,3,4],[3,4,5,6],[6,7,8,9]]
    :return:
        数据按列统计的特征 DataFrame
    """
    statistical_features = {}
    feature_number = feature_values.shape[1]
    for column_idx in range(feature_number):
        column_values = feature_values[:, column_idx]

        # 基本统计特征
        statistical_features[f'max_{column_idx}'] = np.max(column_values)
        statistical_features[f'min_{column_idx}'] = np.min(column_values)
        statistical_features[f'mean_{column_idx}'] = np.mean(column_values)
        statistical_features[f'std_{column_idx}'] = np.std(column_values)

        # 高级统计特征
        statistical_features[f'skew_{column_idx}'] = skew(column_values)
        statistical_features[f'kurtosis_{column_idx}'] = kurtosis(column_values)

        # 变化特征
        statistical_features[f'slope_{column_idx}'] = \
            (np.polyfit(range(len(column_values)), column_values, 1))[0]  # 线性趋势斜率
        statistical_features[f'last_{column_idx}'] = column_values[-1]  # 最新值
        statistical_features[f'diff_{column_idx}'] = column_values[-1] - column_values[0]  # 首尾差值

        # 异常值特征
        q75, q25 = np.percentile(column_values, [75, 25])
        iqr = q75 - q25
        statistical_features[f'outliers_{column_idx}'] = np.sum(
            (column_values < (q25 - 1.5 * iqr)) | (column_values > (q75 + 1.5 * iqr)))

        # 特征两两之间关系的新特征

    return pd.DataFrame([statistical_features])


def extract_window_label(window_labels: NDArray[int], mode='first_failure') -> int:
    """
    一个窗口的标签数据,根据规则计算出一个标签
    :param window_labels: 标签窗口数据一维数组
    :param mode:
        计算故障的模式: first-第一个;last-最后一个;first_failure-第一个故障;last_failure-最后一个故障;
    :return: 故障类型
    """
    if mode == 'first':
        return window_labels[0].item()
    if mode == 'last':
        return window_labels[-1].item()

    indices = np.flatnonzero(window_labels)
    if mode == 'first_failure':
        return window_labels[indices[0]] if indices.size > 0 else 0
    if mode == 'last_failure':
        return window_labels[indices[-1]] if indices.size > 0 else 0
    return 0


def lstm_to_lightgbm(X_lstm):
    """将LSTM输入转换为LightGBM输入"""
    # 展平时间维度
    sample_number = X_lstm.shape[0]
    feature_number = X_lstm.shape[1] * X_lstm.shape[2]

    X_flat = X_lstm.reshape(sample_number, feature_number)
    return X_flat


def lightgbm_to_lstm(X_lightgbm, window_size, sensor_number):
    """将LightGBM输入转换为LSTM输入"""
    # 恢复时间维度
    num_samples = X_lightgbm.shape[0]
    X_lstm = X_lightgbm.reshape(num_samples, window_size, sensor_number)
    return X_lstm


# unit test
if __name__ == "__main__":
    print("start unit test")
    sample_number = 17
    window_size = 6
    datas = np.arange(sample_number)
    print(f'datas: {datas}')
    window_datas = window_data(datas, 12, window_size)
    print(f'window_datas: {window_datas}')

    windowing_datas = windowing_data(datas, window_size)
    print(f'windowing_datas:\n {windowing_datas}')
    window_labels = np.array([0, 0, 1, 2, 3, 2, 0, 3, 2, 0, 0])
    print(f'window_labels:\n {window_labels}')
    window_label = extract_window_label(window_labels, "first")
    print(f'window_label_first:\n {window_label}')
    window_label = extract_window_label(window_labels, "first_failure")
    print(f'window_label_first_failure:\n {window_label}')
    window_label = extract_window_label(window_labels, "last")
    print(f'window_label_last:\n {window_label}')
    window_label = extract_window_label(window_labels, "last_failure")
    print(f'window_label_last_failure:\n {window_label}')

    X_ltsm = np.array([[[0, 1, 2, 3], [4, 5, 6, 7]], [[8, 9, 10, 11], [12, 13, 14, 15]]])
    X_lightgbm = lstm_to_lightgbm(X_ltsm)
    print(f"X_lstm: {X_ltsm} to X_lightgbm: {X_lightgbm}")
    X_ltsm = lightgbm_to_lstm(X_lightgbm, 2, 4)
    print(f"X_lightgbm: {X_lightgbm} to X_lstm: {X_ltsm}")
